import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory

# --- 进阶部分：创建一个窗口历史记录的包装类 ---
class WindowedChatMessageHistory(BaseChatMessageHistory):
    """一个封装了普通历史记录，但只暴露最近 K 个消息的包装类"""
    def __init__(self, base_history: BaseChatMessageHistory, k: int):
        super().__init__()
        self.base_history = base_history
        self.k = k

    @property
    def messages(self):
        # 这里的魔法：当链条来请求历史消息时，我们只返回最后 k*2 条
        # (k*2 是因为一轮对话包含用户和AI两条消息)
        return self.base_history.messages[-(self.k * 2):]

    def add_messages(self, messages) -> None:
        # 将新消息添加到原始的、完整的历史记录中
        self.base_history.add_messages(messages)

    def clear(self) -> None:
        self.base_history.clear()


# --- 环境准备 (与之前相同) ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")
api_key = os.getenv("OPENAI_API_KEY")
model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- 核心代码 ---

    llm = ChatOpenAI(model=model_name)
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个乐于助人的AI助手。"),
        MessagesPlaceholder(variable_name="chat_history"),
        ("user", "{input}"),
    ])
    output_parser = StrOutputParser()
    chain = prompt | llm | output_parser

    # 仍然使用一个字典来模拟数据库
    store = {}

    def get_session_history(session_id: str) -> BaseChatMessageHistory:
        if session_id not in store:
            # 如果是新会话，创建一个完整的历史记录对象
            store[session_id] = ChatMessageHistory()
        
        # 将完整的历史记录对象用我们的窗口包装器包裹起来，窗口大小设为2
        return WindowedChatMessageHistory(base_history=store[session_id], k=2)

    with_message_history = RunnableWithMessageHistory(
        chain,
        get_session_history,
        input_messages_key="input",
        history_messages_key="chat_history",
    )

    session_id = "chat_with_window_memory"
    print("\n--- 开始与 AI 对话 (窗口记忆 k=2) ---")

    # 1. 介绍自己
    print("\n[你]: 你好，我叫小明，我最喜欢的颜色是蓝色。")
    response_1 = with_message_history.invoke(
        {"input": "你好，我叫小明，我最喜欢的颜色是蓝色。"},
        config={"configurable": {"session_id": session_id}}
    )
    print(f"[AI]: {response_1}")

    # 2, 3, 4. 进行一些无关的对话，来"挤掉"第一条记忆
    print("\n[你]: LangChain 是什么？")
    with_message_history.invoke({"input": "LangChain 是什么？"}, config={"configurable": {"session_id": session_id}})
    print("\n[你]: Python 语言有什么特点？")
    with_message_history.invoke({"input": "Python 语言有什么特点？"}, config={"configurable": {"session_id": session_id}})
    
    # 5. 现在来测试它是否还记得你的名字和颜色
    print("\n[你]: 我叫什么名字？我最喜欢的颜色是什么？")
    response_final = with_message_history.invoke(
        {"input": "我叫什么名字？我最喜欢的颜色是什么？"},
        config={"configurable": {"session_id": session_id}}
    )
    print(f"[AI]: {response_final}")
    print("\n--- 注意：如果AI回答不知道，说明窗口记忆起作用了！ ---")

    # 查看完整的、未被包装的记忆
    print(f"\n--- 完整的后台记忆库: ---\n{store[session_id].messages}")
    print("--------------------") 